import torch
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
from typing import Dict, List, Tuple
from sklearn.metrics import mutual_info_score
from scipy.stats import entropy
import seaborn as sns


class ExplainabilityMetrics:
    """可解释性指标计算"""
    
    def __init__(self, save_dir: str = "explainability_results"):
        self.save_dir = Path(save_dir)
        self.save_dir.mkdir(parents=True, exist_ok=True)
    
    def calculate_feature_importance(self, gradients: torch.Tensor, 
                                   input_data: torch.Tensor) -> torch.Tensor:
        """计算特征重要性分数"""
        # 使用梯度×输入的方法
        importance = torch.abs(gradients * input_data)
        return importance.mean(dim=0)  # 对batch维度求平均
    
    def calculate_activation_entropy(self, activations: torch.Tensor) -> float:
        """计算激活熵"""
        # 将激活值转换为概率分布
        flat_activations = activations.flatten().cpu().numpy()
        
        # 归一化到[0,1]
        normalized = (flat_activations - flat_activations.min()) / (flat_activations.max() - flat_activations.min() + 1e-8)
        
        # 计算直方图
        hist, _ = np.histogram(normalized, bins=50, density=True)
        hist = hist + 1e-8  # 避免log(0)
        
        # 计算熵
        return entropy(hist)
    
    def analyze_information_flow(self, model, input1: torch.Tensor, 
                               input2: torch.Tensor) -> Dict[str, float]:
        """分析信息流"""
        activations = {}
        
        def get_activation(name):
            def hook(model, input, output):
                activations[name] = output.detach()
            return hook
        
        # 注册钩子
        hooks = []
        layer_names = []
        for name, module in model.named_modules():
            if isinstance(module, (torch.nn.Conv2d, torch.nn.Linear)):
                hook = module.register_forward_hook(get_activation(name))
                hooks.append(hook)
                layer_names.append(name)
        
        # 前向传播
        with torch.no_grad():
            _ = model(input1, input2)
        
        # 计算每层的信息指标
        info_metrics = {}
        for name in layer_names:
            if name in activations:
                activation = activations[name]
                info_metrics[name] = {
                    'entropy': self.calculate_activation_entropy(activation),
                    'sparsity': (activation == 0).float().mean().item(),
                    'mean_activation': activation.mean().item(),
                    'std_activation': activation.std().item()
                }
        
        # 清理钩子
        for hook in hooks:
            hook.remove()
        
        return info_metrics
    
    def visualize_information_metrics(self, info_metrics: Dict[str, Dict]):
        """可视化信息指标"""
        layer_names = list(info_metrics.keys())
        entropies = [info_metrics[name]['entropy'] for name in layer_names]
        sparsities = [info_metrics[name]['sparsity'] for name in layer_names]
        
        fig, axes = plt.subplots(2, 2, figsize=(15, 10))
        
        # 熵分布
        axes[0, 0].plot(range(len(entropies)), entropies, 'o-')
        axes[0, 0].set_title('Activation Entropy by Layer')
        axes[0, 0].set_xlabel('Layer Index')
        axes[0, 0].set_ylabel('Entropy')
        axes[0, 0].grid(True)
        
        # 稀疏性分布
        axes[0, 1].plot(range(len(sparsities)), sparsities, 'o-', color='orange')
        axes[0, 1].set_title('Activation Sparsity by Layer')
        axes[0, 1].set_xlabel('Layer Index')
        axes[0, 1].set_ylabel('Sparsity Ratio')
        axes[0, 1].grid(True)
        
        # 熵vs稀疏性散点图
        axes[1, 0].scatter(entropies, sparsities, alpha=0.7)
        axes[1, 0].set_xlabel('Entropy')
        axes[1, 0].set_ylabel('Sparsity')
        axes[1, 0].set_title('Entropy vs Sparsity')
        axes[1, 0].grid(True)
        
        # 激活统计
        mean_activations = [info_metrics[name]['mean_activation'] for name in layer_names]
        std_activations = [info_metrics[name]['std_activation'] for name in layer_names]
        
        axes[1, 1].errorbar(range(len(mean_activations)), mean_activations, 
                           yerr=std_activations, fmt='o-', capsize=5)
        axes[1, 1].set_title('Mean Activation ± Std by Layer')
        axes[1, 1].set_xlabel('Layer Index')
        axes[1, 1].set_ylabel('Activation Value')
        axes[1, 1].grid(True)
        
        plt.tight_layout()
        
        save_path = self.save_dir / 'information_metrics.png'
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()
        
        print(f"Information metrics visualization saved to {save_path}")
    
    def compute_layer_similarity(self, activations_dict: Dict[str, torch.Tensor]) -> np.ndarray:
        """计算层间相似性"""
        layer_names = list(activations_dict.keys())
        n_layers = len(layer_names)
        similarity_matrix = np.zeros((n_layers, n_layers))
        
        for i, layer1 in enumerate(layer_names):
            for j, layer2 in enumerate(layer_names):
                if i <= j:
                    # 计算余弦相似度
                    act1 = activations_dict[layer1].flatten()
                    act2 = activations_dict[layer2].flatten()
                    
                    # 确保维度匹配
                    min_len = min(len(act1), len(act2))
                    act1 = act1[:min_len]
                    act2 = act2[:min_len]
                    
                    similarity = torch.cosine_similarity(act1.unsqueeze(0), act2.unsqueeze(0)).item()
                    similarity_matrix[i, j] = similarity
                    similarity_matrix[j, i] = similarity
        
        return similarity_matrix, layer_names
    
    def visualize_layer_similarity(self, similarity_matrix: np.ndarray, 
                                 layer_names: List[str]):
        """可视化层间相似性"""
        plt.figure(figsize=(12, 10))
        
        # 简化层名称显示
        display_names = [name.split('.')[-1] if '.' in name else name for name in layer_names]
        
        sns.heatmap(similarity_matrix, 
                   xticklabels=display_names,
                   yticklabels=display_names,
                   annot=True, 
                   cmap='coolwarm', 
                   center=0,
                   fmt='.3f')
        
        plt.title('Layer Similarity Matrix (Cosine Similarity)')
        plt.xticks(rotation=45, ha='right')
        plt.yticks(rotation=0)
        plt.tight_layout()
        
        save_path = self.save_dir / 'layer_similarity.png'
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()
        
        print(f"Layer similarity visualization saved to {save_path}")
